# LangGraph工作流定义

from langgraph.graph import StateGraph, END
from models.state import WorkflowState
from agents.entity_recognizer import EntityRecognizer
from agents.graph_retriever import GraphRetriever
from agents.intent_classifier import IntentRecognitionAgent
from agents.sql_generator import SQLGenerator
from agents.sql_validator import SQLValidator
from agents.sql_executor import SQLExecutor
from agents.response_generator import ResponseGenerator




def create_workflow():
    # 创建图
    workflow = StateGraph(WorkflowState)
    
    # 初始化各个Agent
    entity_agent = EntityRecognizer()
    retrieval_agent = GraphRetriever()
    intent_agent = IntentRecognitionAgent()
    sql_gen_agent = SQLGenerator()
    sql_val_agent = SQLValidator()
    sql_exec_agent = SQLExecutor()
    response_agent = ResponseGenerator()
    
    # 添加节点
    workflow.add_node("entity_recognition", entity_agent.process)
    workflow.add_node("schema_retrieval", retrieval_agent.process)
    workflow.add_node("intent_recognition", intent_agent.process)
    workflow.add_node("sql_generation", sql_gen_agent.process)
    workflow.add_node("sql_validation", sql_val_agent.process)
    workflow.add_node("sql_execution", sql_exec_agent.process)
    workflow.add_node("response_generation", response_agent.process)
    
    # 设置入口点
    workflow.set_entry_point("entity_recognition")
    
    # 定义边 - 线性流程
    workflow.add_edge("entity_recognition", "schema_retrieval")
    workflow.add_edge("schema_retrieval", "intent_recognition")
    workflow.add_edge("intent_recognition", "sql_generation")
    workflow.add_edge("sql_generation", "sql_validation")
    workflow.add_edge("sql_validation", "sql_execution")
    workflow.add_edge("sql_execution", "response_generation")
    workflow.add_edge("response_generation", END)
    
    return workflow.compile()